{ "cells": [ { "attachments": {}, "cell_type": "markdown", "id": "1d2b5af7", "metadata": {}, "source": [ "Copyright (c) 2023 Graphcore Ltd. All rights reserved." ] }, { "attachments": {}, "cell_type": "markdown", "id": "002700b7", "metadata": {}, "source": [ "# IPU Peak teraFlops\n", "\n", "This notebook shows how you can reach the maximum usages of IPU flops in a couple of lines of Python." ] }, { "attachments": {}, "cell_type": "markdown", "id": "108c181b", "metadata": {}, "source": [ "## Environment setup\n", "\n", "To run the demo using IPU hardware, you need to have the Poplar SDK enabled {and a PopTorch/TensorFlow wheel installed}. Refer to the [Getting Started guide](https://docs.graphcore.ai/en/latest/getting-started.html#getting-started) for your system for details on how to do this. Also refer to the [Jupyter Quick Start guide](https://docs.graphcore.ai/projects/jupyter-notebook-quick-start/en/latest/index.html) for how to set up Jupyter to be able to run this notebook on a remote IPU machine." ] }, { "attachments": {}, "cell_type": "markdown", "id": "5f856b38", "metadata": {}, "source": [ "## Dependencies" ] }, { "attachments": {}, "cell_type": "markdown", "id": "f14b0107", "metadata": {}, "source": [ "Import dependencies and define configuration. " ] }, { "cell_type": "code", "execution_count": 80, "id": "17226e0c", "metadata": {}, "outputs": [], "source": [ "from jax.config import config\n", "\n", "# Select how many IPUs will be visible.\n", "config.FLAGS.jax_ipu_device_count = 1" ] }, { "cell_type": "code", "execution_count": 2, "id": "3d8870fc", "metadata": {}, "outputs": [], "source": [ "import jax\n", "import jax.lax\n", "import numpy as np\n", "import plotly.graph_objects as go\n", "from tqdm import tqdm" ] }, { "cell_type": "code", "execution_count": null, "id": "fdc49153", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 3, "id": "174d33ec", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "IpuDevice(id=0, num_tiles=1472, version=ipu2)" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Need real IPU hardware!\n", "d = jax.devices(\"ipu\")[0]\n", "assert not d.is_ipu_model\n", "d" ] }, { "cell_type": "code", "execution_count": 4, "id": "bf6e85a4", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "1850000000.0" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "num_tiles = d.num_tiles\n", "tiles = tuple(range(num_tiles))\n", "d.tile_clock_frequency" ] }, { "cell_type": "code", "execution_count": null, "id": "08afe55c", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 5, "id": "86825ab8", "metadata": {}, "outputs": [], "source": [ "from tessellate_ipu.tile import IpuConvVertexType, ipu_cycle_count, ipu_cycle_count_overhead, tile_map, tile_put_sharded\n", "\n", "\n", "def tile_basic_matmul(lhs, rhs):\n", " \"\"\"Run a basic matmul on every tile, using the `ConvPartial1x1` vertex.\n", "\n", " Most optimal hardware usage: loading a small weight matrix once in CCCS\n", " registers, and then performing a \"long\" AMP pipeline of size N.\n", "\n", " Args:\n", " lhs: [N, 8/16] array\n", " rhs: [8/16, 8/16] array.\n", " Returns:\n", " Matmul result. [N, 8/16]\n", " \"\"\"\n", " accumulator_dtype = lhs.dtype\n", " # accumulator_dtype = np.float32\n", " output = tile_map(\n", " jax.lax.dot_general_p,\n", " lhs,\n", " rhs,\n", " dimension_numbers=(([1], [1]), ([], [])),\n", " precision=jax.lax.Precision.DEFAULT,\n", " preferred_element_type=accumulator_dtype,\n", " ipu_vertex_type=IpuConvVertexType.ConvPartial1x1,\n", " )\n", " return output\n", "\n", "\n", "def tile_benchmark_matmul(lhs, rhs):\n", " \"\"\"Benchmarking matmul IPU.\"\"\"\n", " lhs = tile_put_sharded(lhs, tiles)\n", " rhs = tile_put_sharded(rhs, tiles)\n", " # Get IPU raw cycle before and after matmul\n", " lhs, rhs, start = ipu_cycle_count(lhs, rhs)\n", " out = tile_basic_matmul(lhs, rhs)\n", " out, end = ipu_cycle_count(out)\n", " return out, start, end\n", "\n", "\n", "# CPU jitting to double check the result!\n", "tile_benchmark_matmul_ipu = jax.jit(tile_benchmark_matmul, backend=\"ipu\")\n", "tile_benchmark_matmul_cpu = jax.jit(tile_benchmark_matmul, backend=\"cpu\")" ] }, { "cell_type": "code", "execution_count": 6, "id": "3e2b5bfc", "metadata": {}, "outputs": [], "source": [ "def measure_tile_matmul_flops(N, dtype):\n", " \"\"\"Measure IPU flops usage by running small matmuls on every tile.\"\"\"\n", " dtype = np.float16\n", " lhs_size = N\n", " rhs_size = 16\n", " contract_size = 16\n", "\n", " lhs_data = np.random.randn(len(tiles), lhs_size, contract_size).astype(dtype)\n", " rhs_data = np.random.randn(len(tiles), rhs_size, contract_size).astype(dtype)\n", " # Run independent matmuls on every tile, and measure cycle count.\n", " _, start, end = tile_benchmark_matmul_ipu(lhs_data, rhs_data)\n", " tile_cycle_count = np.asarray(end.array) - np.asarray(start.array)\n", " cycle_count = np.mean(tile_cycle_count[:, 0]) - ipu_cycle_count_overhead()\n", " # Flops conversion.\n", " execution_time = cycle_count / d.tile_clock_frequency\n", " # Count product & sum in matmul as floating point operation.\n", " num_ops = 2 * lhs_size * rhs_size * contract_size * num_tiles\n", " tflops = num_ops / execution_time * 1e-12\n", " return cycle_count, tflops" ] }, { "cell_type": "code", "execution_count": 52, "id": "6d0851f4", "metadata": {}, "outputs": [], "source": [ "# Matrix sizes to check!\n", "max_size = 1024 * 2\n", "step_size = 16\n", "sizes = np.arange(16, max_size + step_size, step_size)" ] }, { "cell_type": "code", "execution_count": 53, "id": "1b5acf11", "metadata": {}, "outputs": [], "source": [ "# IPU Mk2 peak tflops.\n", "peak_tflops = d.tile_clock_frequency * num_tiles * 128 * 1e-12\n", "peak_tflops = np.array([peak_tflops for _ in sizes])" ] }, { "cell_type": "code", "execution_count": 54, "id": "338d72b9", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 128/128 [09:17<00:00, 4.36s/it]\n" ] } ], "source": [ "cycle_counts = []\n", "tflops = []\n", "\n", "for s in tqdm(sizes):\n", " c, f = measure_tile_matmul_flops(s, np.float16)\n", " cycle_counts.append(c)\n", " tflops.append(f)\n", "\n", "cycle_counts = np.asarray(cycle_counts)\n", "tflops = np.asarray(tflops)" ] }, { "cell_type": "code", "execution_count": 55, "id": "1cc693ac", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "(array([ 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176,\n", " 192, 208, 224, 240, 256, 272, 288, 304, 320, 336, 352,\n", " 368, 384, 400, 416, 432, 448, 464, 480, 496, 512, 528,\n", " 544, 560, 576, 592, 608, 624, 640, 656, 672, 688, 704,\n", " 720, 736, 752, 768, 784, 800, 816, 832, 848, 864, 880,\n", " 896, 912, 928, 944, 960, 976, 992, 1008, 1024, 1040, 1056,\n", " 1072, 1088, 1104, 1120, 1136, 1152, 1168, 1184, 1200, 1216, 1232,\n", " 1248, 1264, 1280, 1296, 1312, 1328, 1344, 1360, 1376, 1392, 1408,\n", " 1424, 1440, 1456, 1472, 1488, 1504, 1520, 1536, 1552, 1568, 1584,\n", " 1600, 1616, 1632, 1648, 1664, 1680, 1696, 1712, 1728, 1744, 1760,\n", " 1776, 1792, 1808, 1824, 1840, 1856, 1872, 1888, 1904, 1920, 1936,\n", " 1952, 1968, 1984, 2000, 2016, 2032, 2048]),\n", " array([ 55.49366766, 96.70991447, 127.72015878, 152.79763288,\n", " 170.72418021, 186.94235531, 201.23605773, 211.11875736,\n", " 221.11904141, 230.45923967, 236.55806879, 243.36495709,\n", " 250.00854069, 254.0509482 , 258.99908359, 264.00537751,\n", " 266.82080765, 270.58772183, 274.52113575, 276.54860334,\n", " 279.52120668, 282.71082765, 284.20808654, 286.61825782,\n", " 289.26937759, 290.39578638, 292.3923635 , 294.63996377,\n", " 295.49748827, 297.18189698, 299.11854948, 299.77583828,\n", " 301.21890147, 302.91032332, 303.4163886 , 304.66781426,\n", " 306.16202255, 306.55200957, 307.64841641, 308.98136288,\n", " 309.27992413, 310.25002808, 311.44920104, 311.67522347,\n", " 312.5406127 , 313.62741516, 313.79524575, 314.5728 ,\n", " 315.56416443, 315.68389506, 316.38797953, 317.29749147,\n", " 317.37924575, 318.01914931, 318.85784615, 318.90841331,\n", " 319.4929399 , 320.2698899 , 320.29384353, 320.83107958,\n", " 321.5538087 , 321.5563996 , 322.05147278, 322.72628427,\n", " 322.7102802 , 323.16900579, 323.80122288, 323.76958892,\n", " 324.19615703, 324.79030948, 324.74585655, 325.14346494,\n", " 325.70343424, 325.64817259, 326.01989088, 326.54902435,\n", " 326.48463306, 326.83310353, 327.3343049 , 327.26284615,\n", " 327.58970384, 328.06550588, 327.98718028, 328.29540463,\n", " 328.74802774, 328.66417002, 328.95517505, 329.38657503,\n", " 329.29774649, 329.57335785, 329.98526502, 329.89250389,\n", " 330.15376499, 330.54771652, 330.4522102 , 330.69975639,\n", " 331.07712313, 330.97829062, 331.21430498, 331.57631391,\n", " 331.47544484, 331.70005085, 332.04780393, 331.94476957,\n", " 332.15934657, 332.49383667, 332.3897526 , 332.59429531,\n", " 332.91641972, 332.81082649, 333.006783 , 333.31735496,\n", " 333.21052703, 333.39850572, 333.69826431, 333.5909087 ,\n", " 333.77099294, 334.06061157, 333.95245959, 334.12562756,\n", " 334.40572131, 334.29718004, 334.46366299, 334.73479497,\n", " 334.62645228, 334.78623802, 335.04892488, 334.94046716]))" ] }, "execution_count": 55, "metadata": {}, "output_type": "execute_result" } ], "source": [ "sizes, tflops" ] }, { "cell_type": "code", "execution_count": null, "id": "d9573c7e", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 77, "id": "366d9584", "metadata": {}, "outputs": [ { "data": { "application/vnd.plotly.v1+json": { "config": { "plotlyServerURL": "https://plot.ly" }, "data": [ { "line": { "dash": "dot" }, "mode": "lines", "name": "IPU Mk2 Peak TFlops", "showlegend": true, "type": "scatter", "x": [ 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512, 528, 544, 560, 576, 592, 608, 624, 640, 656, 672, 688, 704, 720, 736, 752, 768, 784, 800, 816, 832, 848, 864, 880, 896, 912, 928, 944, 960, 976, 992, 1008, 1024, 1040, 1056, 1072, 1088, 1104, 1120, 1136, 1152, 1168, 1184, 1200, 1216, 1232, 1248, 1264, 1280, 1296, 1312, 1328, 1344, 1360, 1376, 1392, 1408, 1424, 1440, 1456, 1472, 1488, 1504, 1520, 1536, 1552, 1568, 1584, 1600, 1616, 1632, 1648, 1664, 1680, 1696, 1712, 1728, 1744, 1760, 1776, 1792, 1808, 1824, 1840, 1856, 1872, 1888, 1904, 1920, 1936, 1952, 1968, 1984, 2000, 2016, 2032, 2048 ], "y": [ 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696, 348.5696 ] }, { "marker": { "size": 5 }, "mode": "markers", "name": "FP16[N,16] @ FP16[16,16] -> FP16", "showlegend": true, "type": "scatter", "x": [ 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512, 528, 544, 560, 576, 592, 608, 624, 640, 656, 672, 688, 704, 720, 736, 752, 768, 784, 800, 816, 832, 848, 864, 880, 896, 912, 928, 944, 960, 976, 992, 1008, 1024, 1040, 1056, 1072, 1088, 1104, 1120, 1136, 1152, 1168, 1184, 1200, 1216, 1232, 1248, 1264, 1280, 1296, 1312, 1328, 1344, 1360, 1376, 1392, 1408, 1424, 1440, 1456, 1472, 1488, 1504, 1520, 1536, 1552, 1568, 1584, 1600, 1616, 1632, 1648, 1664, 1680, 1696, 1712, 1728, 1744, 1760, 1776, 1792, 1808, 1824, 1840, 1856, 1872, 1888, 1904, 1920, 1936, 1952, 1968, 1984, 2000, 2016, 2032, 2048 ], "y": [ 55.49366766169154, 96.70991446612008, 127.72015877862594, 152.79763287671233, 170.72418020895722, 186.94235530726255, 201.23605773195877, 211.11875736484444, 221.11904140969162, 230.45923966942146, 236.55806878745966, 243.36495709090912, 250.00854068965518, 254.05094820270466, 258.99908359133127, 264.0053775147929, 266.82080765351185, 270.58772183288414, 274.52113575129533, 276.54860333629773, 279.52120668257754, 282.71082764976956, 284.2080865448065, 286.6182578158458, 289.269377593361, 290.3957863835368, 292.3923634951456, 294.63996377358495, 295.4974882729825, 297.18189698046183, 299.11854948096885, 299.7758382793814, 301.2189014729951, 302.91032332268367, 303.41638859884773, 304.66781426403645, 306.16202255192877, 306.5520095687498, 307.648416407355, 308.9813628808864, 309.27992412694266, 310.25002807947016, 311.44920103896106, 311.67522347211894, 312.5406127023661, 313.6274151589242, 313.79524574573804, 314.5728, 315.5641644341801, 315.6838950559652, 316.38797953281426, 317.2974914660831, 317.37924575003206, 318.019149313622, 318.85784615384614, 318.90841330926713, 319.4929398994975, 320.2698899009901, 320.2938435253663, 320.83107957814, 321.5538086956522, 321.55639960123676, 322.05147277726854, 322.7262842676311, 322.7102802016507, 323.1690057945566, 323.8012228769497, 323.7695889164313, 324.1961570345409, 324.790309484193, 324.74585655040545, 325.14346493927127, 325.70343424, 325.6481725869757, 326.0198908807482, 326.5490243451464, 326.4846330649351, 326.8331035311796, 327.33430490341755, 327.2628461484032, 327.58970384336476, 328.0655058823529, 327.9871802838813, 328.2954046250876, 328.748027739251, 328.66417002361106, 328.9551750508475, 329.38657503355705, 329.2977464905837, 329.57335784635586, 329.9852650195058, 329.8925038879945, 330.1537649904519, 330.547716519546, 330.4522102037617, 330.69975639283507, 331.07712313341494, 330.9782906163891, 331.21430497900417, 331.5763139120095, 331.4754448422978, 331.70005084548103, 332.0478039306358, 331.9447695667556, 332.15934656834935, 332.49383667041616, 332.38975260256996, 332.5942953064605, 332.9164197152245, 332.810826490238, 333.0067830016138, 333.3173549626467, 333.21052702702895, 333.3985057157839, 333.6982643080125, 333.5909086991856, 333.7709929411764, 334.06061157360404, 333.95245958783903, 334.125627558662, 334.4057213082259, 334.2971800377701, 334.4636629936616, 334.7347949661181, 334.6264522755953, 334.7862380181038, 335.04892488174073, 334.9404671550231 ] } ], "layout": { "font": { "family": "Courier New" }, "template": { "data": { "bar": [ { "error_x": { "color": "#2a3f5f" }, "error_y": { "color": "#2a3f5f" }, "marker": { "line": { "color": "#E5ECF6", "width": 0.5 }, "pattern": { "fillmode": "overlay", "size": 10, "solidity": 0.2 } }, "type": "bar" } ], "barpolar": [ { "marker": { "line": { "color": "#E5ECF6", "width": 0.5 }, "pattern": { "fillmode": "overlay", "size": 10, "solidity": 0.2 } }, "type": "barpolar" } ], "carpet": [ { "aaxis": { "endlinecolor": "#2a3f5f", "gridcolor": "white", "linecolor": "white", "minorgridcolor": "white", "startlinecolor": "#2a3f5f" }, "baxis": { "endlinecolor": "#2a3f5f", "gridcolor": "white", "linecolor": "white", "minorgridcolor": "white", "startlinecolor": "#2a3f5f" }, "type": "carpet" } ], "choropleth": [ { "colorbar": { "outlinewidth": 0, "ticks": "" }, "type": "choropleth" } ], "contour": [ { "colorbar": { "outlinewidth": 0, "ticks": "" }, "colorscale": [ [ 0, "#0d0887" ], [ 0.1111111111111111, "#46039f" ], [ 0.2222222222222222, "#7201a8" ], [ 0.3333333333333333, "#9c179e" ], [ 0.4444444444444444, "#bd3786" ], [ 0.5555555555555556, "#d8576b" ], [ 0.6666666666666666, "#ed7953" ], [ 0.7777777777777778, "#fb9f3a" ], [ 0.8888888888888888, "#fdca26" ], [ 1, "#f0f921" ] ], "type": "contour" } ], "contourcarpet": [ { "colorbar": { "outlinewidth": 0, "ticks": "" }, "type": "contourcarpet" } ], "heatmap": [ { "colorbar": { "outlinewidth": 0, "ticks": "" }, "colorscale": [ [ 0, "#0d0887" ], [ 0.1111111111111111, "#46039f" ], [ 0.2222222222222222, "#7201a8" ], [ 0.3333333333333333, "#9c179e" ], [ 0.4444444444444444, "#bd3786" ], [ 0.5555555555555556, "#d8576b" ], [ 0.6666666666666666, "#ed7953" ], [ 0.7777777777777778, "#fb9f3a" ], [ 0.8888888888888888, "#fdca26" ], [ 1, "#f0f921" ] ], "type": "heatmap" } ], "heatmapgl": [ { "colorbar": { "outlinewidth": 0, "ticks": "" }, "colorscale": [ [ 0, "#0d0887" ], [ 0.1111111111111111, "#46039f" ], [ 0.2222222222222222, "#7201a8" ], [ 0.3333333333333333, "#9c179e" ], [ 0.4444444444444444, "#bd3786" ], [ 0.5555555555555556, "#d8576b" ], [ 0.6666666666666666, "#ed7953" ], [ 0.7777777777777778, "#fb9f3a" ], [ 0.8888888888888888, "#fdca26" ], [ 1, "#f0f921" ] ], "type": "heatmapgl" } ], "histogram": [ { "marker": { "pattern": { "fillmode": "overlay", "size": 10, "solidity": 0.2 } }, "type": "histogram" } ], "histogram2d": [ { "colorbar": { "outlinewidth": 0, "ticks": "" }, "colorscale": [ [ 0, "#0d0887" ], [ 0.1111111111111111, "#46039f" ], [ 0.2222222222222222, "#7201a8" ], [ 0.3333333333333333, "#9c179e" ], [ 0.4444444444444444, "#bd3786" ], [ 0.5555555555555556, "#d8576b" ], [ 0.6666666666666666, "#ed7953" ], [ 0.7777777777777778, "#fb9f3a" ], [ 0.8888888888888888, "#fdca26" ], [ 1, "#f0f921" ] ], "type": "histogram2d" } ], "histogram2dcontour": [ { "colorbar": { "outlinewidth": 0, "ticks": "" }, "colorscale": [ [ 0, "#0d0887" ], [ 0.1111111111111111, "#46039f" ], [ 0.2222222222222222, "#7201a8" ], [ 0.3333333333333333, "#9c179e" ], [ 0.4444444444444444, "#bd3786" ], [ 0.5555555555555556, "#d8576b" ], [ 0.6666666666666666, "#ed7953" ], [ 0.7777777777777778, "#fb9f3a" ], [ 0.8888888888888888, "#fdca26" ], [ 1, "#f0f921" ] ], "type": "histogram2dcontour" } ], "mesh3d": [ { "colorbar": { "outlinewidth": 0, "ticks": "" }, "type": "mesh3d" } ], "parcoords": [ { "line": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "type": "parcoords" } ], "pie": [ { "automargin": true, "type": "pie" } ], "scatter": [ { "fillpattern": { "fillmode": "overlay", "size": 10, "solidity": 0.2 }, "type": "scatter" } ], "scatter3d": [ { "line": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "marker": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "type": "scatter3d" } ], "scattercarpet": [ { "marker": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "type": "scattercarpet" } ], "scattergeo": [ { "marker": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "type": "scattergeo" } ], "scattergl": [ { "marker": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "type": "scattergl" } ], "scattermapbox": [ { "marker": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "type": "scattermapbox" } ], "scatterpolar": [ { "marker": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "type": "scatterpolar" } ], "scatterpolargl": [ { "marker": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "type": "scatterpolargl" } ], "scatterternary": [ { "marker": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "type": "scatterternary" } ], "surface": [ { "colorbar": { "outlinewidth": 0, "ticks": "" }, "colorscale": [ [ 0, "#0d0887" ], [ 0.1111111111111111, "#46039f" ], [ 0.2222222222222222, "#7201a8" ], [ 0.3333333333333333, "#9c179e" ], [ 0.4444444444444444, "#bd3786" ], [ 0.5555555555555556, "#d8576b" ], [ 0.6666666666666666, "#ed7953" ], [ 0.7777777777777778, "#fb9f3a" ], [ 0.8888888888888888, "#fdca26" ], [ 1, "#f0f921" ] ], "type": "surface" } ], "table": [ { "cells": { "fill": { "color": "#EBF0F8" }, "line": { "color": "white" } }, "header": { "fill": { "color": "#C8D4E3" }, "line": { "color": "white" } }, "type": "table" } ] }, "layout": { "annotationdefaults": { "arrowcolor": "#2a3f5f", "arrowhead": 0, "arrowwidth": 1 }, "autotypenumbers": "strict", "coloraxis": { "colorbar": { "outlinewidth": 0, "ticks": "" } }, "colorscale": { "diverging": [ [ 0, "#8e0152" ], [ 0.1, "#c51b7d" ], [ 0.2, "#de77ae" ], [ 0.3, "#f1b6da" ], [ 0.4, "#fde0ef" ], [ 0.5, "#f7f7f7" ], [ 0.6, "#e6f5d0" ], [ 0.7, "#b8e186" ], [ 0.8, "#7fbc41" ], [ 0.9, "#4d9221" ], [ 1, "#276419" ] ], "sequential": [ [ 0, "#0d0887" ], [ 0.1111111111111111, "#46039f" ], [ 0.2222222222222222, "#7201a8" ], [ 0.3333333333333333, "#9c179e" ], [ 0.4444444444444444, "#bd3786" ], [ 0.5555555555555556, "#d8576b" ], [ 0.6666666666666666, "#ed7953" ], [ 0.7777777777777778, "#fb9f3a" ], [ 0.8888888888888888, "#fdca26" ], [ 1, "#f0f921" ] ], "sequentialminus": [ [ 0, "#0d0887" ], [ 0.1111111111111111, "#46039f" ], [ 0.2222222222222222, "#7201a8" ], [ 0.3333333333333333, "#9c179e" ], [ 0.4444444444444444, "#bd3786" ], [ 0.5555555555555556, "#d8576b" ], [ 0.6666666666666666, "#ed7953" ], [ 0.7777777777777778, "#fb9f3a" ], [ 0.8888888888888888, "#fdca26" ], [ 1, "#f0f921" ] ] }, "colorway": [ "#636efa", "#EF553B", "#00cc96", "#ab63fa", "#FFA15A", "#19d3f3", "#FF6692", "#B6E880", "#FF97FF", "#FECB52" ], "font": { "color": "#2a3f5f" }, "geo": { "bgcolor": "white", "lakecolor": "white", "landcolor": "#E5ECF6", "showlakes": true, "showland": true, "subunitcolor": "white" }, "hoverlabel": { "align": "left" }, "hovermode": "closest", "mapbox": { "style": "light" }, "paper_bgcolor": "white", "plot_bgcolor": "#E5ECF6", "polar": { "angularaxis": { "gridcolor": "white", "linecolor": "white", "ticks": "" }, "bgcolor": "#E5ECF6", "radialaxis": { "gridcolor": "white", "linecolor": "white", "ticks": "" } }, "scene": { "xaxis": { "backgroundcolor": "#E5ECF6", "gridcolor": "white", "gridwidth": 2, "linecolor": "white", "showbackground": true, "ticks": "", "zerolinecolor": "white" }, "yaxis": { "backgroundcolor": "#E5ECF6", "gridcolor": "white", "gridwidth": 2, "linecolor": "white", "showbackground": true, "ticks": "", "zerolinecolor": "white" }, "zaxis": { "backgroundcolor": "#E5ECF6", "gridcolor": "white", "gridwidth": 2, "linecolor": "white", "showbackground": true, "ticks": "", "zerolinecolor": "white" } }, "shapedefaults": { "line": { "color": "#2a3f5f" } }, "ternary": { "aaxis": { "gridcolor": "white", "linecolor": "white", "ticks": "" }, "baxis": { "gridcolor": "white", "linecolor": "white", "ticks": "" }, "bgcolor": "#E5ECF6", "caxis": { "gridcolor": "white", "linecolor": "white", "ticks": "" } }, "title": { "x": 0.05 }, "xaxis": { "automargin": true, "gridcolor": "white", "linecolor": "white", "ticks": "", "title": { "standoff": 15 }, "zerolinecolor": "white", "zerolinewidth": 2 }, "yaxis": { "automargin": true, "gridcolor": "white", "linecolor": "white", "ticks": "", "title": { "standoff": 15 }, "zerolinecolor": "white", "zerolinewidth": 2 } } }, "title": { "text": "IPU on-tile AMP micro-benchmarking (TFLOPS)" }, "xaxis": { "tickvals": [ 16, 64, 128, 256, 512, 768, 1024, 1536, 2048 ], "title": { "text": "N (input shape: [N,16])" } }, "yaxis": { "title": { "text": "TFLOPS" } } } }, "text/html": [ "
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "fig = go.Figure()\n", "\n", "# Peak TFlops.\n", "fig.add_trace(\n", " go.Scatter(\n", " x=sizes,\n", " y=peak_tflops,\n", " mode=\"lines\",\n", " name=\"IPU Mk2 Peak TFlops\",\n", " line=go.scatter.Line(dash=\"dot\"),\n", " showlegend=True,\n", " )\n", ")\n", "# FP16 TFlops.\n", "fig.add_trace(\n", " go.Scatter(\n", " x=sizes,\n", " y=tflops,\n", " mode=\"markers\",\n", " name=\"FP16[N,16] @ FP16[16,16] -> FP16\",\n", " marker=go.scatter.Marker(size=5),\n", " showlegend=True,\n", " )\n", ")\n", "\n", "fig.update_layout(\n", " font_family=\"Courier New\",\n", " title_text=\"IPU on-tile AMP micro-benchmarking (TFLOPS)\",\n", " xaxis_title=\"N (input shape: [N,16])\",\n", " yaxis_title=\"TFLOPS\",\n", ")\n", "fig.update_xaxes(tickvals=[16, 64, 128, 256, 512, 768, 1024, 1536, 2048])\n", "fig.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "3c08ad92", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.10" } }, "nbformat": 4, "nbformat_minor": 5 }